{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\".\")\n",
    "sys.path.append(\"../../\")\n",
    "\n",
    "import os\n",
    "os.chdir(\"../..\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wiss/wimbauer/miniconda3/envs/campred/lib/python3.11/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.\n",
      "  @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)\n",
      "/home/wiss/wimbauer/miniconda3/envs/campred/lib/python3.11/site-packages/ignite/handlers/checkpoint.py:16: DeprecationWarning: `TorchScript` support for functional optimizers is deprecated and will be removed in a future PyTorch release. Consider using the `torch.compile` optimizer instead.\n",
      "  from torch.distributed.optim import ZeroRedundancyOptimizer\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from moviepy.video.io.ImageSequenceClip import ImageSequenceClip\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import copy\n",
    "\n",
    "import hydra\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from einops import einsum\n",
    "\n",
    "from hydra import compose, initialize\n",
    "from omegaconf import OmegaConf\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from anycam.common.array_operations import map_fn, unsqueezer\n",
    "from anycam.visualization.common import color_tensor\n",
    "\n",
    "from torch.cuda.amp import autocast\n",
    "import math\n",
    "\n",
    "import uuid\n",
    "\n",
    "from torchvision.utils import flow_to_image\n",
    "\n",
    "from anycam.datasets import make_datasets\n",
    "from anycam.models import make_depth_predictor, make_pose_predictor\n",
    "from anycam.trainer import CamPredWrapper\n",
    "from anycam.loss.metric import camera_to_rel_deg\n",
    "from anycam.datasets.common import get_flow_selector, get_index_selector\n",
    "\n",
    "\n",
    "from anycam.datasets.common import get_ids_for_sequence\n",
    "\n",
    "from anycam.scripts.fit_video import fit_video\n",
    "\n",
    "import rerun as rr\n",
    "import rerun.blueprint as rrb\n",
    "\n",
    "\n",
    "from minipytorch3d.rotation_conversions import (\n",
    "    matrix_to_quaternion,\n",
    "    quaternion_to_matrix,\n",
    "    matrix_to_axis_angle,\n",
    "    axis_angle_to_matrix,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_depth(depth, threshold=0.1):\n",
    "    _, h, w = depth.shape\n",
    "\n",
    "    depth = depth.clone()[None, ...]\n",
    "    median = torch.median(depth)\n",
    "    \n",
    "    depth_grad = torch.stack(torch.gradient(depth, dim=(-2, -1))).norm(dim=0)\n",
    "\n",
    "    mask = depth_grad < median * threshold\n",
    "\n",
    "    return mask\n",
    "\n",
    "\n",
    "def visualize_rerun(trajectory, depths, imgs, proj, uncertainties=None, cam_every=1, depth_every=1, subsample_pts=16, step=0, uncertainty_thresh=-1, radii=1, track_masks=None, max_depth_masks=None, filter_depth_threshold=0.1, static_points_accumulate=False, highlight_dynamic=False, image_plane_distance=0.05):\n",
    "    h, w = imgs[0].shape[:2]\n",
    "\n",
    "    rr.set_time_sequence(\"step\", step)\n",
    "\n",
    "    if cam_every > 0:\n",
    "        for i, d_id in enumerate(list(range(len(depths)))[cam_every//2::cam_every]):\n",
    "            if static_points_accumulate:\n",
    "                rr.set_time_sequence(\"step\", d_id)\n",
    "\n",
    "            pose = trajectory[d_id]\n",
    "            rr.log(f\"world/scene/cameras/{i:03d}\", rr.Pinhole(\n",
    "                resolution=[w, h],\n",
    "                focal_length=float(proj[0, 0]),\n",
    "                image_plane_distance=image_plane_distance, \n",
    "            ), static=not static_points_accumulate)\n",
    "            rr.log(f\"world/scene/cameras/{i:03d}\", rr.Transform3D(translation=pose[:3, 3].cpu(), mat3x3=pose[:3, :3].cpu(), axis_length=0.01), static=not static_points_accumulate)\n",
    "\n",
    "            depth_img = color_tensor(depths[d_id], cmap=\"plasma\")[0]\n",
    "            depth_img = depth_img.cpu().numpy()\n",
    "            depth_img = (depth_img * 255).astype(np.uint8)\n",
    "\n",
    "            # rr.log(f\"world/scene/cameras/{i:03d}/depth\", rr.Image(depth_img).compress())\n",
    "\n",
    "            if uncertainties is not None:\n",
    "                uncertainty_img = color_tensor(uncertainties[d_id].clamp(max=torch.quantile(uncertainties[d_id], .95)), cmap=\"plasma\", norm=True)[0]\n",
    "                uncertainty_img = uncertainty_img.cpu().numpy()\n",
    "                uncertainty_img = (uncertainty_img * 255).astype(np.uint8)\n",
    "\n",
    "                # rr.log(f\"world/scene/cameras/{i:03d}/uncertainty\", rr.Image(uncertainty_img).compress())\n",
    "\n",
    "            rr.log(f\"world/scene/cameras/{i:03d}/input\", rr.Image((imgs[d_id] * 255).astype(np.uint8)).compress(jpeg_quality=95), static=not static_points_accumulate)\n",
    "\n",
    "            if highlight_dynamic > 0 and (d_id - cam_every//2) % cam_every == 0:\n",
    "                pts, colors = lift_image(torch.tensor(imgs[d_id]).cuda(), depths[d_id].cuda(), trajectory[d_id].cuda(), proj)\n",
    "                dyn_mask = uncertainties[d_id].view(-1) > highlight_dynamic\n",
    "                dyn_mask = dyn_mask & filter_depth(depths[d_id].cuda(), threshold=filter_depth_threshold).view(-1)\n",
    "\n",
    "                pts_dyn = pts[dyn_mask, :]\n",
    "                colors_dyn = colors[dyn_mask, :]\n",
    "\n",
    "                colors_dyn = colors_dyn * .5 + colors_dyn.new_tensor((1, 0, 0)) * .5\n",
    "\n",
    "                colors_dyn = (colors_dyn * 255).clamp(0, 255).to(torch.uint8)\n",
    "\n",
    "                rr.log(f\"world/scene/points/{i:03d}_dyn\", rr.Points3D(pts_dyn[:, :3].cpu().numpy(), colors=colors_dyn[:, :3].cpu().numpy(), radii=rr.Radius.ui_points([radii]),), static=not static_points_accumulate)\n",
    "\n",
    "\n",
    "    xyz_min, xyz_max = None, None\n",
    "\n",
    "    for i, t_id in enumerate(list(range(len(depths)))[depth_every//2::depth_every]):\n",
    "        pts, colors = lift_image(torch.tensor(imgs[t_id]).cuda(), depths[t_id].cuda(), trajectory[t_id].cuda(), proj)\n",
    "        mask = filter_depth(depths[t_id].cuda(), threshold=filter_depth_threshold)\n",
    "\n",
    "        mask = mask.view(-1)\n",
    "\n",
    "        mask = mask & (depths[t_id].cuda().view(-1) < 10)\n",
    "\n",
    "        if uncertainty_thresh > 0:\n",
    "            mask = mask & (uncertainties[t_id].view(-1) < uncertainty_thresh)\n",
    "\n",
    "        if track_masks is not None:\n",
    "            mask = mask & track_masks[t_id].view(-1)\n",
    "\n",
    "        if max_depth_masks is not None:\n",
    "            mask = mask & max_depth_masks[t_id].view(-1)           \n",
    "\n",
    "        pts_filtered = pts[mask, :]\n",
    "        colors_filtered = colors[mask, :]\n",
    "\n",
    "        pts_filtered = pts_filtered[subsample_pts//2::subsample_pts]\n",
    "        colors_filtered = colors_filtered[subsample_pts//2::subsample_pts]\n",
    "\n",
    "        if static_points_accumulate:\n",
    "            rr.set_time_sequence(\"step\", t_id)\n",
    "\n",
    "        colors_filtered = (colors_filtered * 255).clamp(0, 255).to(torch.uint8)\n",
    "\n",
    "        rr.log(f\"world/scene/points/{i:03d}\", rr.Points3D(pts_filtered[:, :3].cpu().numpy(), colors=colors_filtered[:, :3].cpu().numpy(), radii=rr.Radius.ui_points([radii]),), static=not static_points_accumulate)\n",
    "    \n",
    "\n",
    "        if xyz_min is None and pts_filtered.shape[0] > 0:\n",
    "            xyz_min = pts_filtered.min(dim=0).values\n",
    "            xyz_max = pts_filtered.max(dim=0).values\n",
    "        elif pts_filtered.shape[0] > 0:\n",
    "            xyz_min = torch.min(xyz_min, pts_filtered.min(dim=0).values)\n",
    "            xyz_max = torch.max(xyz_max, pts_filtered.max(dim=0).values)\n",
    "\n",
    "    return xyz_min, xyz_max\n",
    "\n",
    "\n",
    "def visualize_video_rerun(\n",
    "        trajectory, \n",
    "        depths, \n",
    "        imgs, \n",
    "        proj,\n",
    "        uncertainties=None, \n",
    "        cam_every=-1, \n",
    "        depth_every=1, \n",
    "        subsample_pts=1, \n",
    "        subsample_dynamic_pts=1, \n",
    "        step=0, \n",
    "        radii=1.5, \n",
    "        uncertainty_thresh=-1, \n",
    "        track_masks=None, \n",
    "        max_depth_masks=None, \n",
    "        follow_cam=None, \n",
    "        filter_depth_threshold=0.1,\n",
    "        static_points_accumulate=False,\n",
    "        highlight_dynamic=False,\n",
    "        image_plane_distance=0.05,\n",
    "        ):\n",
    "    h, w = imgs[0].shape[:2]\n",
    "\n",
    "    imgs = np.array(imgs)\n",
    "\n",
    "    # imgs_grey = (imgs.astype(float) * .5 + .5)\n",
    "    imgs_grey = imgs\n",
    "\n",
    "    xyz_min, xyz_max = visualize_rerun(\n",
    "        trajectory, \n",
    "        depths, \n",
    "        imgs_grey, \n",
    "        proj, \n",
    "        uncertainties, \n",
    "        cam_every=cam_every, \n",
    "        depth_every=depth_every, \n",
    "        subsample_pts=subsample_pts, \n",
    "        step=step, \n",
    "        uncertainty_thresh=uncertainty_thresh, \n",
    "        radii=radii, \n",
    "        track_masks=track_masks, \n",
    "        max_depth_masks=max_depth_masks, \n",
    "        filter_depth_threshold=filter_depth_threshold,\n",
    "        static_points_accumulate=static_points_accumulate,\n",
    "        highlight_dynamic=highlight_dynamic,\n",
    "        image_plane_distance=image_plane_distance,\n",
    "        )\n",
    "\n",
    "    for t_id in range(len(depths)):\n",
    "        rr.set_time_sequence(\"step\", t_id+step)\n",
    "\n",
    "        pts, colors = lift_image(torch.tensor(imgs[t_id]).cuda(), depths[t_id].cuda() - 0.01, trajectory[t_id].cuda(), proj)\n",
    "        mask = filter_depth(depths[t_id].cuda(), threshold=filter_depth_threshold)\n",
    "\n",
    "        mask = mask.view(-1)\n",
    "\n",
    "        mask = mask & (uncertainties[t_id].view(-1) >= uncertainty_thresh)\n",
    "\n",
    "        if max_depth_masks is not None:\n",
    "            mask = mask & max_depth_masks[t_id].view(-1)\n",
    "\n",
    "        pts = pts[mask, :]\n",
    "        colors = colors[mask, :]\n",
    "\n",
    "        pts = pts[subsample_dynamic_pts//2::subsample_dynamic_pts]\n",
    "        colors = colors[subsample_dynamic_pts//2::subsample_dynamic_pts]\n",
    "        colors = (colors * 255).clamp(0, 255).to(torch.uint8)\n",
    "\n",
    "        if xyz_min is None and pts.shape[0] > 0:\n",
    "            xyz_min = pts.min(dim=0).values\n",
    "            xyz_max = pts.max(dim=0).values\n",
    "        elif pts.shape[0] > 0:\n",
    "            xyz_min = torch.min(xyz_min, pts.min(dim=0).values)\n",
    "            xyz_max = torch.max(xyz_max, pts.max(dim=0).values)\n",
    "\n",
    "        rr.log(f\"world/scene/active_points\", rr.Points3D(pts[:, :3].cpu().numpy(), colors=colors[:, :3].cpu().numpy(), radii=rr.Radius.ui_points([radii]),))\n",
    "\n",
    "        pose = trajectory[t_id]\n",
    "        rot = quaternion_to_matrix(matrix_to_quaternion(pose[:3, :3].cpu()))\n",
    "\n",
    "        rr.log(f\"world/scene/active_cam\", rr.Pinhole(\n",
    "            resolution=[w, h],\n",
    "            focal_length=float(proj[0, 0]),\n",
    "            image_plane_distance=image_plane_distance, \n",
    "        ), static=True)\n",
    "        rr.log(f\"world/scene/active_cam\", rr.Transform3D(translation=pose[:3, 3].cpu(), mat3x3=rot, axis_length=0.01))\n",
    "\n",
    "        rr.log(\"world/scene/cam_traj\", rr.LineStrips3D([pose[:3, 3].cpu().numpy().tolist() for pose in trajectory[:t_id+1]], colors=[(0, 255, 0)]), static=False)\n",
    "\n",
    "        rr.log(\"world/scene/active_cam/input\", rr.Image((imgs[t_id] * 255).astype(np.uint8)).compress(jpeg_quality=95))\n",
    "        \n",
    "        if uncertainties is not None:\n",
    "            uncertainty_img = color_tensor((uncertainties[t_id] / 0.05).clamp(0, 1), cmap=\"plasma\", norm=False)[0]\n",
    "            uncertainty_img = uncertainty_img.cpu().numpy()\n",
    "            uncertainty_img = (uncertainty_img * 255).astype(np.uint8)\n",
    "\n",
    "            rr.log(f\"world/scene/active_cam/uncertainty\", rr.Image(uncertainty_img).compress())\n",
    "\n",
    "\n",
    "        if follow_cam == \"world\":\n",
    "            rr.log(\"world/scene\", rr.Transform3D(translation=pose[:3, 3].cpu(), mat3x3=rot, axis_length=0, from_parent=True))\n",
    "\n",
    "        elif follow_cam is not None:\n",
    "            fc_t, fc_a = follow_cam\n",
    "            rr.log(f\"world/scene/ac/follow_cam\", rr.Pinhole(\n",
    "                resolution=[1280, 720],\n",
    "                focal_length=1280,\n",
    "                image_plane_distance=0.02, \n",
    "            ), static=True)\n",
    "            rr.log(f\"world/scene/ac\", rr.Transform3D(translation=pose[:3, 3].cpu(), mat3x3=rot, axis_length=0.01))\n",
    "            rr.log(f\"world/scene/ac/follow_cam\", rr.Transform3D(translation=fc_t, rotation=rr.RotationAxisAngle(axis=[0, 1, 0], angle=rr.Angle(deg=fc_a)) ,axis_length=0.01))\n",
    "\n",
    "    xyz_abs_max = torch.max(torch.abs(xyz_min), torch.abs(xyz_max)).cpu()\n",
    "    bbox_pts = torch.tensor([\n",
    "        [-1, -1, -1],\n",
    "        [1, -1, -1],\n",
    "        [1, 1, -1],\n",
    "        [-1, 1, -1],\n",
    "        [-1, -1, 1],\n",
    "        [1, -1, 1],\n",
    "        [1, 1, 1],\n",
    "        [-1, 1, 1],\n",
    "    ]) * xyz_abs_max\n",
    "\n",
    "    # rr.log(\"world/scene/bbox\", rr.Points3D(bbox_pts.numpy(), colors=[(255, 255, 255)]), static=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from anycam.common.geometry import get_grid_xy\n",
    "\n",
    "\n",
    "def get_normals(depth, proj):\n",
    "    # Get grid\n",
    "\n",
    "    n, _, h, w = depth.shape\n",
    "\n",
    "    xy = get_grid_xy(h, w, homogeneous=True, device=depth.device)\n",
    "\n",
    "    # Project to 3D\n",
    "    pts = torch.inverse(proj) @ xy.reshape(3, -1)\n",
    "\n",
    "    pts = pts * depth.view(n, 1, -1)\n",
    "\n",
    "    # Get normals\n",
    "    pts = pts.reshape(n, 3, h, w)\n",
    "\n",
    "    dx, dy = torch.gradient(pts, dim=(-2, -1))\n",
    "\n",
    "    n = torch.cross(dx, dy, dim=1)\n",
    "\n",
    "    n = n / n.norm(dim=1, keepdim=True)\n",
    "\n",
    "    return n\n",
    "\n",
    "\n",
    "def lift_image(img, depth, pose, proj):\n",
    "    h, w = img.shape[:2]\n",
    "\n",
    "    proj = torch.tensor(proj, device=device).float()\n",
    "\n",
    "    proj[0, 0] = proj[0, 0] / w * 2\n",
    "    proj[1, 1] = proj[1, 1] / h * 2\n",
    "    proj[0, 2] = proj[0, 2] / w * 2 - 1\n",
    "    proj[1, 2] = proj[1, 2] / h * 2 - 1\n",
    "\n",
    "    inv_proj = torch.inverse(proj)\n",
    "\n",
    "    pts = get_grid_xy(h, w, homogeneous=True).reshape(3, h*w).cuda()\n",
    "    pts = inv_proj @ pts\n",
    "    pts = pts * depth.view(1, -1).cuda()\n",
    "    pts = torch.cat((pts, torch.ones(1, h*w, device=device)), dim=0)\n",
    "    pts = pose @ pts\n",
    "    pts = pts[:3, :].T\n",
    "\n",
    "    colors = torch.tensor(img.reshape(-1, 3)).cuda()\n",
    "\n",
    "    return pts, colors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_depth_flow(imgs, model):\n",
    "    depths = []\n",
    "    flow_occs_fwd = []\n",
    "    flow_occs_bwd = []\n",
    "\n",
    "    c, h, w = imgs.shape[1:]\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for i in tqdm(range(len(imgs)-1)):\n",
    "            img0 = imgs[i]\n",
    "            img1 = imgs[i+1]\n",
    "\n",
    "            imgs_ = torch.stack([img0, img1]).unsqueeze(0).cuda()\n",
    "\n",
    "            images_ip_fwd, images_ip_bwd = model.image_processor(imgs_ * 2 - 1, data={})\n",
    "\n",
    "            flow_occs_fwd.append(images_ip_fwd[0, :(1 if i != len(imgs)-2 else 2), 3:6].cpu())\n",
    "            flow_occs_bwd.append(images_ip_bwd[0, (1 if i != 0 else 0):, 3:6].cpu())\n",
    "\n",
    "            depth = model.depth_predictor(img0.unsqueeze(0).cuda())\n",
    "\n",
    "            depth = 1 / depth[0].clamp_min(1e-3)\n",
    "\n",
    "            depths.append(depth.cpu())\n",
    "\n",
    "        depth = model.depth_predictor(imgs[-1].unsqueeze(0).cuda())\n",
    "        depth = 1 / depth[0].clamp_min(1e-3)\n",
    "\n",
    "        depths.append(depth.cpu())\n",
    "\n",
    "    seq_depths = torch.cat(depths, dim=0)\n",
    "    seq_flow_occs_fwd = torch.cat(flow_occs_fwd, dim=0)\n",
    "    seq_flow_occs_bwd = torch.cat(flow_occs_bwd, dim=0)\n",
    "\n",
    "    return seq_depths, seq_flow_occs_fwd, seq_flow_occs_bwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from depthcrafter.depth_crafter_ppl import DepthCrafterPipeline\n",
    "from depthcrafter.unet import DiffusersUNetSpatioTemporalConditionModelDepthCrafter\n",
    "\n",
    "\n",
    "def get_depth_crafter():\n",
    "    unet = DiffusersUNetSpatioTemporalConditionModelDepthCrafter.from_pretrained(\n",
    "        \"tencent/DepthCrafter\",\n",
    "        low_cpu_mem_usage=True,\n",
    "        torch_dtype=torch.float16,\n",
    "    )\n",
    "    pipe = DepthCrafterPipeline.from_pretrained(\n",
    "        \"stabilityai/stable-video-diffusion-img2vid-xt\",\n",
    "        unet=unet,\n",
    "        torch_dtype=torch.float16,\n",
    "        variant=\"fp16\",\n",
    "    )\n",
    "    pipe.to(\"cuda\")\n",
    "    return pipe\n",
    "\n",
    "def depth_crafter_inference(pipe, imgs):\n",
    "    # imgs = imgs.permute(0, 2, 3, 1).float().cuda()\n",
    "    n, h, w, c = imgs.shape\n",
    "\n",
    "    imgs = imgs.permute(0, 3, 1, 2)\n",
    "\n",
    "    th = math.ceil(h / 128) * 128\n",
    "    tw = math.ceil(w / 128) * 128\n",
    "\n",
    "    imgs = F.interpolate(imgs, (th, tw))\n",
    "\n",
    "    imgs = imgs.permute(0, 2, 3, 1).cpu().numpy()\n",
    "\n",
    "    res = pipe(\n",
    "            imgs,\n",
    "            height=imgs.shape[1],\n",
    "            width=imgs.shape[2],\n",
    "            output_type=\"np\",\n",
    "            guidance_scale=1,\n",
    "            num_inference_steps=5,\n",
    "            window_size=110,\n",
    "            overlap=25,\n",
    "            track_time=True,\n",
    "        ).frames[0]\n",
    "\n",
    "    res = res.sum(-1) / res.shape[-1]\n",
    "\n",
    "    res = F.interpolate(torch.tensor(res, device=\"cuda\").unsqueeze(1), (h, w)).cpu()\n",
    "\n",
    "    res = 1 / res\n",
    "\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def depth_alignment(depths_uni, depths_dc, uncertainties):\n",
    "    n, _, h, w = depths_uni.shape\n",
    "\n",
    "    depths_uni_opt = depths_uni[:-1].reshape(-1).cuda()\n",
    "    depths_dc_opt = depths_dc[:-1].reshape(-1).cuda()\n",
    "    uncertainties = uncertainties.reshape(-1)\n",
    "\n",
    "    mask = (depths_uni_opt > 0) & (depths_uni_opt < 50)\n",
    "\n",
    "    depths_uni_opt = depths_uni_opt[mask]\n",
    "    depths_dc_opt = depths_dc_opt[mask]\n",
    "    uncertainties = uncertainties[mask]\n",
    "\n",
    "    depths_uni_opt = 1 / depths_uni_opt\n",
    "    depths_dc_opt = 1 / depths_dc_opt\n",
    "\n",
    "    A = torch.stack([depths_dc_opt, torch.ones_like(depths_dc_opt)], dim=1)\n",
    "    b = depths_uni_opt.view(-1, 1)\n",
    "\n",
    "    weight = 1 / uncertainties.clamp_min(1e-3)\n",
    "    weight = weight.view(-1, 1)\n",
    "\n",
    "    A = A * weight\n",
    "    b = b * weight\n",
    "\n",
    "    x = torch.linalg.lstsq(A, b).solution\n",
    "\n",
    "    scale = x[0]\n",
    "    shift = x[1]\n",
    "\n",
    "    print(f\"Aligning with scale {scale.item()} and shift {shift.item()}\")\n",
    "\n",
    "    depths_dc_aligned = 1 / ((1 / depths_dc.cuda()) * scale + shift).clamp_min(1e-3)\n",
    "\n",
    "    return depths_dc_aligned.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_track_masks(flow_occs, uncertainties, max_track_len=8, uncertainty_thresh=0.1, min_track_len=1, plot_rerun=False, imgs=None):\n",
    "    n, _, h, w = flow_occs.shape\n",
    "    device = flow_occs.device\n",
    "\n",
    "    masks = torch.zeros_like(flow_occs[:, :1, :, :])\n",
    "    tracks = torch.zeros((0, 2), device=device, dtype=torch.long)\n",
    "    track_lens = torch.zeros((0,), device=device, dtype=torch.long)\n",
    "    track_starts = torch.zeros((0, 3), device=device, dtype=torch.float)\n",
    "\n",
    "    \n",
    "    if plot_rerun:\n",
    "        rr.init(\"Track Masks\", recording_id=uuid.uuid4())\n",
    "        rr.connect()\n",
    "\n",
    "    for i_curr_frame in range(n):\n",
    "        if plot_rerun:\n",
    "            rr.set_time_sequence(\"step\", i_curr_frame)\n",
    "            rr.log(\"2d/img\", rr.Image((imgs[i_curr_frame] * 255).astype(np.uint8)).compress())\n",
    "\n",
    "        # If maximum track length has been reached, clear the track\n",
    "        len_mask = track_lens < max_track_len\n",
    "\n",
    "        tracks = tracks[len_mask]\n",
    "        track_lens = track_lens[len_mask]\n",
    "        track_starts = track_starts[len_mask]\n",
    "\n",
    "\n",
    "        # Create mask of pixels that are currently being tracked\n",
    "\n",
    "        tracked_pixels = torch.zeros_like(masks[0, 0], dtype=torch.bool)\n",
    "\n",
    "        tracked_pixels[tracks[:, 1], tracks[:, 0]] = True\n",
    "\n",
    "        \n",
    "        # Add new tracks for untracked pixels\n",
    "\n",
    "        x = torch.arange(w, device=device).view(1, -1).expand(h, -1)\n",
    "        y = torch.arange(h, device=device).view(-1, 1).expand(-1, w)\n",
    "        xy = torch.stack([x, y], dim=-1).view(-1, 2)\n",
    "\n",
    "        new_tracks = xy[~tracked_pixels.view(-1), :]\n",
    "        new_track_lens = torch.ones_like(new_tracks[:, 0])\n",
    "        new_track_starts = torch.cat([new_tracks, i_curr_frame * torch.ones_like(new_tracks[:, 0]).view(-1, 1)], dim=1)\n",
    "\n",
    "\n",
    "        # Filter both tracks and new tracks by uncertainty\n",
    "\n",
    "        track_uncs = uncertainties[i_curr_frame, 0, tracks[:, 1], tracks[:, 0]]\n",
    "        new_track_uncs = uncertainties[i_curr_frame, 0, new_tracks[:, 1], new_tracks[:, 0]]\n",
    "\n",
    "        tracks = tracks[track_uncs < uncertainty_thresh]\n",
    "        track_lens = track_lens[track_uncs < uncertainty_thresh]\n",
    "        track_starts = track_starts[track_uncs < uncertainty_thresh]\n",
    "\n",
    "        new_tracks = new_tracks[new_track_uncs < uncertainty_thresh]\n",
    "        new_track_lens = new_track_lens[new_track_uncs < uncertainty_thresh]\n",
    "        new_track_starts = new_track_starts[new_track_uncs < uncertainty_thresh]\n",
    "\n",
    "        if plot_rerun:\n",
    "            rr.log(\"2d/tracks\", rr.Points2D(tracks.cpu().numpy(), colors=[(255, 255, 0)]))\n",
    "            rr.log(\"2d/new_tracks\", rr.Points2D(new_tracks.cpu().numpy(), colors=[(0, 255, 0)]))\n",
    "\n",
    "\n",
    "        # Add new tracks to track list\n",
    "        tracks = torch.cat([tracks, new_tracks], dim=0)\n",
    "        track_lens = torch.cat([track_lens, new_track_lens], dim=0)\n",
    "        track_starts = torch.cat([track_starts, new_track_starts], dim=0)\n",
    "\n",
    "\n",
    "        # Register new tracks in mask\n",
    "\n",
    "        # masks[i_curr_frame, 0, new_tracks[:, 1], new_tracks[:, 0]] = True\n",
    "        \n",
    "        register_mask = track_lens == min_track_len\n",
    "\n",
    "        masks[i_curr_frame, 0, tracks[register_mask, 1], tracks[register_mask, 0]] = True\n",
    "\n",
    "\n",
    "        # Check if track is occluded. If so, remove it\n",
    "\n",
    "        track_occs = flow_occs[i_curr_frame, 2, :, :][tracks[:, 1], tracks[:, 0]] > .5\n",
    "\n",
    "\n",
    "        if plot_rerun:\n",
    "            rr.log(\"2d/occluded_tracks\", rr.Points2D(tracks[~track_occs].cpu().numpy(), colors=[(255, 0, 0)]))\n",
    "\n",
    "\n",
    "        tracks = tracks[track_occs]\n",
    "        track_lens = track_lens[track_occs]\n",
    "        track_starts = track_starts[track_occs]\n",
    "\n",
    "\n",
    "        # Compute flow for tracked pixels\n",
    "\n",
    "        track_flow_x = flow_occs[i_curr_frame, 0, :, :][tracks[:, 1], tracks[:, 0]]\n",
    "        track_flow_y = flow_occs[i_curr_frame, 1, :, :][tracks[:, 1], tracks[:, 0]]\n",
    "\n",
    "        track_flow_x = track_flow_x * w\n",
    "        track_flow_y = track_flow_y * h\n",
    "\n",
    "        track_flow = torch.stack([track_flow_x, track_flow_y], dim=-1)\n",
    "\n",
    "        tracks = tracks.float() + track_flow\n",
    "        tracks = tracks.round().long()\n",
    "\n",
    "\n",
    "        # Check if track is out of bounds. If so, remove it\n",
    "\n",
    "        out_of_bounds = (tracks[:, 0] < 0) | (tracks[:, 0] >= w) | (tracks[:, 1] < 0) | (tracks[:, 1] >= h)\n",
    "\n",
    "        tracks = tracks[~out_of_bounds]\n",
    "        track_lens = track_lens[~out_of_bounds]\n",
    "        track_starts = track_starts[~out_of_bounds]\n",
    "\n",
    "        \n",
    "        # Update track_lens\n",
    "\n",
    "        track_lens += 1\n",
    "\n",
    "\n",
    "    return masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from anycam.datasets.davis.davis_dataset import DavisDataset\n",
    "\n",
    "\n",
    "davis = DavisDataset(\n",
    "    data_path=\"data/Davis/2017-trainval\",\n",
    "    split_path=None,\n",
    "    image_size=[336, int(16/9*336)],\n",
    "    # image_size=336,\n",
    "    frame_count=2,\n",
    "    return_depth=False,\n",
    "    return_flow=False,\n",
    "    dilation=1,\n",
    "    index_selector=get_index_selector(True),\n",
    "    flow_selector=get_flow_selector(1, True)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from anycam.datasets.common import get_sequence_sampler\n",
    "from anycam.datasets.waymo.waymo_dataset import WaymoDataset\n",
    "\n",
    "waymo =  WaymoDataset(\n",
    "    data_path=\"data/waymo/testing\",\n",
    "    split_path=\"anycam/datasets/waymo/splits/eval_seqs_2_64/test_files.txt\",\n",
    "    image_size=[336, int(16/9*336)],\n",
    "    frame_count=2,\n",
    "    return_depth=False,\n",
    "    return_flow=False,\n",
    "    index_selector=get_index_selector(True),\n",
    "    sequence_sampler=get_sequence_sampler(True),\n",
    "    flow_selector=get_flow_selector(0, True),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 30/30 [00:00<00:00, 78349.39it/s]\n"
     ]
    }
   ],
   "source": [
    "from anycam.datasets.aria_everyday_activities.aea_dataset import ExtractedAEADataset\n",
    "from anycam.datasets.common import get_index_selector\n",
    "\n",
    "\n",
    "ext_aea = ExtractedAEADataset(\n",
    "    data_path=\"data/aea_extracted\",\n",
    "    split_path=None,\n",
    "    image_size=336,\n",
    "    frame_count=3,\n",
    "    return_depth=False,\n",
    "    return_flow=False,\n",
    "    dilation=1,\n",
    "    index_selector=get_index_selector(True),\n",
    "    flow_selector=get_flow_selector(1, True),\n",
    "    preprocessed_path=\"/storage/slurm/wimbauer/unimatch_flows/aea/\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 30/30 [00:00<00:00, 102801.57it/s]\n"
     ]
    }
   ],
   "source": [
    "ext_aea = ExtractedAEADataset(\n",
    "    data_path=\"anycam/aea_extracted\",\n",
    "    split_path=None,\n",
    "    image_size=336,\n",
    "    frame_count=3,\n",
    "    return_depth=False,\n",
    "    return_flow=False,\n",
    "    dilation=1,\n",
    "    index_selector=get_index_selector(True),\n",
    "    flow_selector=get_flow_selector(1, True),\n",
    "    # preprocessed_path=\"data/unimatch_flows/aea/\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "\n",
    "def load_and_resize_frames(folder_path, target_size=336):\n",
    "    frames = []\n",
    "    for filename in sorted(os.listdir(folder_path)):\n",
    "        if filename.endswith(('.png', '.jpg', '.jpeg')):\n",
    "            file_path = os.path.join(folder_path, filename)\n",
    "            frame = cv2.imread(file_path)\n",
    "            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "           \n",
    "            height, width = frame.shape[:2]\n",
    "            if height < width:\n",
    "                new_height = target_size\n",
    "                new_width = int((target_size / height) * width)\n",
    "            else:\n",
    "                new_width = target_size\n",
    "                new_height = int((target_size / width) * height)\n",
    "            \n",
    "            resized_frame = cv2.resize(frame, (new_width, new_height))\n",
    "\n",
    "            resized_frame = resized_frame.astype(np.float32) / 255.0\n",
    "\n",
    "            frames.append(resized_frame)\n",
    "    \n",
    "    return frames\n",
    "\n",
    "folder_path = 'custom_data/running_dog_pexels'\n",
    "imgs = load_and_resize_frames(folder_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'train': 80,\n",
       " 'swing': 60,\n",
       " 'bmx-bumps': 90,\n",
       " 'drift-chicane': 52,\n",
       " 'varanus-cage': 67,\n",
       " 'horsejump-low': 60,\n",
       " 'dog-agility': 25,\n",
       " 'bus': 80,\n",
       " 'motocross-jump': 40,\n",
       " 'parkour': 100,\n",
       " 'soapbox': 99,\n",
       " 'tractor-sand': 76,\n",
       " 'drone': 91,\n",
       " 'lucia': 70,\n",
       " 'upside-down': 65,\n",
       " 'paragliding': 70,\n",
       " 'rallye': 50,\n",
       " 'kite-walk': 80,\n",
       " 'stroller': 91,\n",
       " 'lady-running': 65,\n",
       " 'blackswan': 50,\n",
       " 'schoolgirls': 80,\n",
       " 'scooter-board': 91,\n",
       " 'bear': 82,\n",
       " 'flamingo': 80,\n",
       " 'bmx-trees': 80,\n",
       " 'drift-straight': 50,\n",
       " 'libby': 49,\n",
       " 'disc-jockey': 76,\n",
       " 'elephant': 80,\n",
       " 'pigs': 79,\n",
       " 'crossing': 52,\n",
       " 'judo': 34,\n",
       " 'classic-car': 63,\n",
       " 'boxing-fisheye': 87,\n",
       " 'bike-packing': 69,\n",
       " 'cows': 104,\n",
       " 'dancing': 62,\n",
       " 'tuk-tuk': 59,\n",
       " 'car-shadow': 40,\n",
       " 'rhino': 90,\n",
       " 'breakdance': 84,\n",
       " 'dance-jump': 60,\n",
       " 'dog-gooses': 86,\n",
       " 'india': 81,\n",
       " 'koala': 100,\n",
       " 'kite-surf': 50,\n",
       " 'night-race': 46,\n",
       " 'lindy-hop': 73,\n",
       " 'car-turn': 80,\n",
       " 'surf': 55,\n",
       " 'breakdance-flare': 71,\n",
       " 'camel': 90,\n",
       " 'dance-twirl': 90,\n",
       " 'mallard-water': 80,\n",
       " 'mallard-fly': 70,\n",
       " 'kid-football': 68,\n",
       " 'hockey': 75,\n",
       " 'walking': 72,\n",
       " 'rollerblade': 35,\n",
       " 'shooting': 40,\n",
       " 'motorbike': 43,\n",
       " 'color-run': 84,\n",
       " 'drift-turn': 64,\n",
       " 'car-roundabout': 75,\n",
       " 'miami-surf': 70,\n",
       " 'planes-water': 38,\n",
       " 'boat': 75,\n",
       " 'scooter-gray': 75,\n",
       " 'mbike-trick': 79,\n",
       " 'dog': 60,\n",
       " 'skate-park': 80,\n",
       " 'motocross-bumps': 60,\n",
       " 'goat': 90,\n",
       " 'tennis': 70,\n",
       " 'dogs-jump': 66,\n",
       " 'soccerball': 48,\n",
       " 'sheep': 68,\n",
       " 'stunt': 71,\n",
       " 'scooter-black': 43,\n",
       " 'paragliding-launch': 80,\n",
       " 'cat-girl': 89,\n",
       " 'horsejump-high': 50,\n",
       " 'longboard': 52,\n",
       " 'lab-coat': 47,\n",
       " 'snowboard': 66,\n",
       " 'dogs-scale': 83,\n",
       " 'gold-fish': 78,\n",
       " 'loading': 50,\n",
       " 'hike': 80}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "davis._sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 63 datapoints\n"
     ]
    }
   ],
   "source": [
    "from anycam.scripts.common import load_model\n",
    "from anycam.loss import make_loss\n",
    "from anycam.trainer import AnyCamWrapper\n",
    "\n",
    "def load_anycam(model_path, checkpoint=None):\n",
    "    with initialize(version_base=None, config_path=str(\"../..\" / model_path), job_name=\"test_app\"):\n",
    "        config = compose(config_name=\"training_config\")\n",
    "\n",
    "    prefix = \"training_checkpoint_\"\n",
    "    ckpts = Path(model_path).glob(f\"{prefix}*.pt\")\n",
    "\n",
    "    model_conf = config[\"model\"]\n",
    "    model_conf[\"use_provided_flow\"] = False\n",
    "    model_conf[\"train_directions\"] = \"forward\"\n",
    "\n",
    "    model = AnyCamWrapper(model_conf)\n",
    "\n",
    "    criterion = [make_loss(cfg) for cfg in config.get(\"loss\", [])][0]\n",
    "\n",
    "    training_steps = [int(ckpt.stem.split(prefix)[1]) for ckpt in ckpts]\n",
    "\n",
    "    print(training_steps)\n",
    "    if training_steps:\n",
    "        if checkpoint is None:\n",
    "            ckpt_path = f\"{prefix}{max(training_steps)}.pt\"\n",
    "        else:\n",
    "            ckpt_path = checkpoint\n",
    "\n",
    "        # ckpt_path = Path(config[\"output\"][\"path\"]) / ckpt_path\n",
    "        ckpt_path = Path(model_path) / ckpt_path\n",
    "\n",
    "        print(ckpt_path)\n",
    "\n",
    "        cp = torch.load(ckpt_path, map_location=\"cpu\")\n",
    "\n",
    "        cp[\"model\"] = {k: v for k, v in cp[\"model\"].items() if not \"depth_predictor\" in k}\n",
    "\n",
    "        model.load_state_dict(cp[\"model\"], strict=False)\n",
    "\n",
    "    return model, criterion\n",
    "\n",
    "\n",
    "run_name = \"anycam_seq8\"\n",
    "\n",
    "# model_path = Path(\"out/vacation_runs\") / run_name\n",
    "model_path = Path(\"pretrained_models\") / run_name\n",
    "\n",
    "device = \"cuda\"\n",
    "\n",
    "dataset = davis\n",
    "sequence = \"drift-turn\"\n",
    "\n",
    "# ids = list(range(515, 575))\n",
    "# ids = list(range(896, 960)) + list(range(1920, 1984))\n",
    "ids, _ = get_ids_for_sequence(dataset, sequence)\n",
    "# ids = [i for i, dp in enumerate(dataset._datapoints) if dataset._datapoints[0][0] == dp[0]]\n",
    "\n",
    "# ids = ids[80:280]\n",
    "\n",
    "print(f\"Found {len(ids)} datapoints\")\n",
    "\n",
    "# imgs = []\n",
    "# poses = []\n",
    "# proj = None\n",
    "\n",
    "# gt_proj = None\n",
    "\n",
    "# for id in ids:\n",
    "#     data = dataset[id]\n",
    "#     imgs.append(data[\"imgs\"][0].transpose(1, 2, 0))\n",
    "#     poses.append(data[\"poses\"][0])\n",
    "#     proj = data[\"projs\"][0]\n",
    "\n",
    "out_path = Path(\"media\") / \"anycam\" / f\"{dataset.NAME.strip()}_{run_name}\"\n",
    "out_path.mkdir(exist_ok=True, parents=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /home/wiss/wimbauer/.cache/torch/hub/lpiccinelli-eth_UniDepth_main\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pytorch3D is not available. Either install it or compile knn under unidepth/ops/knn with `bash compile.sh`\n",
      "Not loading pretrained weights for backbone\n",
      "EdgeGuidedLocalSSI reverts to a non cuda-optimized operation, you will experince large slowdown, please install it:  `cd ./unidepth/ops/extract_patches && bash compile.sh`\n",
      "UniDepth_v2_vits14 is loaded with:\n",
      "\t missing keys: []\n",
      "\t additional keys: []\n",
      "[245000, 247500, 242500, 240000]\n",
      "out/post_submission/anycam_baseline_fc8_drop_backend-nccl-2_5051794/training_checkpoint_247500.pt\n"
     ]
    }
   ],
   "source": [
    "model, criterion = load_anycam(model_path)\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotdict import dotdict\n",
    "\n",
    "fit_video_config = {\n",
    "    \"with_rerun\": False,\n",
    "    \"do_ba_refinement\": True,\n",
    "    \"overfit\": False,\n",
    "    \"prediction\": {\n",
    "        \"model_seq_len\": 100,\n",
    "        \"shift\": 99,\n",
    "        # \"model_seq_len\": 8,\n",
    "        # \"shift\": 7,\n",
    "        \"square_crop\": True,\n",
    "        \"return_all_uncerts\": True,\n",
    "    },\n",
    "    \"ba_refinement\": {\n",
    "        \"with_rerun\": True,\n",
    "        \"max_uncert\": 0.05,\n",
    "        \"lambda_smoothness\": 0.1,\n",
    "        # \"lambda_smoothness\": 0.001,\n",
    "        \"long_tracks\": True,\n",
    "        \"n_steps_last_global\": 20000,\n",
    "    },\n",
    "    \"ba_refinement_level\": 2,\n",
    "    # \"ba_refinement_level\": 0,\n",
    "    \"dataset\": {\n",
    "        \"image_size\": [336, None]\n",
    "        # \"image_size\": [None, 336],\n",
    "    }\n",
    "}\n",
    "\n",
    "fit_video_config = dotdict(fit_video_config)\n",
    "\n",
    "best_trajectory, proj, extras_dict, ba_extras = fit_video(\n",
    "    fit_video_config,\n",
    "    model,\n",
    "    criterion,\n",
    "    imgs[::1],\n",
    "    return_extras=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2210566/1449744676.py:1: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  best_trajectory = torch.tensor(best_trajectory)\n",
      "/tmp/ipykernel_2210566/1449744676.py:4: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  proj_rel = torch.tensor(proj).clone()\n",
      "/tmp/ipykernel_2210566/1449744676.py:13: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  proj = torch.tensor(proj)\n",
      "/tmp/ipykernel_2210566/1449744676.py:14: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  proj_rel = torch.tensor(proj_rel).cuda()\n"
     ]
    }
   ],
   "source": [
    "best_trajectory = torch.tensor(best_trajectory)\n",
    "\n",
    "\n",
    "proj_rel = torch.tensor(proj).clone()\n",
    "\n",
    "h, w = imgs[0].shape[:2]\n",
    "\n",
    "proj_rel[0, 0] = proj_rel[0, 0] / w * 2\n",
    "proj_rel[0, 2] = proj_rel[0, 2] / w * 2 - 1\n",
    "proj_rel[1, 1] = proj_rel[1, 1] / h * 2\n",
    "proj_rel[1, 2] = proj_rel[1, 2] / h * 2 - 1\n",
    "\n",
    "proj = torch.tensor(proj)\n",
    "proj_rel = torch.tensor(proj_rel).cuda()\n",
    "\n",
    "best_candidate = extras_dict[\"best_candidate\"]\n",
    "\n",
    "depths = extras_dict[\"seq_depths\"]\n",
    "uncertainties = torch.stack(extras_dict[\"uncertainties\"])[:, 0, best_candidate, :1, :, :]\n",
    "\n",
    "flows_occs_fwd = extras_dict[\"seq_flow_occs_fwd\"].cuda()\n",
    "flows_occs_bwd = extras_dict[\"seq_flow_occs_bwd\"].cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 134/134 [00:35<00:00,  3.72it/s]\n"
     ]
    }
   ],
   "source": [
    "# Get UniDepth depths and flows\n",
    "\n",
    "depths_uni, seq_flows_occs_fwd, seq_flows_occs_bwd = get_depth_flow(torch.tensor(np.array(imgs), device=\"cuda\").permute(0, 3, 1, 2), model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install `accelerate` for faster and less memory-intense model loading. You can do so with: \n",
      "```\n",
      "pip install accelerate\n",
      "```\n",
      ".\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install `accelerate` for faster and less memory-intense model loading. You can do so with: \n",
      "```\n",
      "pip install accelerate\n",
      "```\n",
      ".\n",
      "Expected types for unet: ['UNetSpatioTemporalConditionModel'], got DiffusersUNetSpatioTemporalConditionModelDepthCrafter.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d1b5e2fe1414f2e874b49b2f835ed75",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time for encoding video: 13927.62109375 ms\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "43c90794f1fc465192670acc3ebfd70a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "81a6640a9528431a9e985e846304e905",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time for denoising video: 24647.189453125 ms\n",
      "Elapsed time for decoding video: 22097.677734375 ms\n"
     ]
    }
   ],
   "source": [
    "# Get DepthCrafter depths\n",
    "\n",
    "dc_pipe = get_depth_crafter()\n",
    "\n",
    "depths_dc = depth_crafter_inference(dc_pipe, torch.tensor(np.array(imgs), device=\"cuda\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([134, 1, 189, 336])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "uncertainties.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Aligning with scale 1.9591450691223145 and shift 0.043867602944374084\n"
     ]
    }
   ],
   "source": [
    "depths_dc_uni = depth_alignment(depths_uni, depths_dc, uncertainties)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_depth_masks = (depths_uni < 1.5) & (depths_dc_uni < 2)\n",
    "max_depth_masks = (depths_uni < 5) & (depths_dc_uni < 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "uncertainties_1 = torch.cat((uncertainties, uncertainties[-1:]), dim=0)\n",
    "uncert_thresh = 0.01\n",
    "\n",
    "track_masks = compute_track_masks(seq_flows_occs_fwd.cuda(), uncertainties_1.cuda(), max_track_len=16, min_track_len=2, uncertainty_thresh=uncert_thresh, plot_rerun=False, imgs=imgs)\n",
    "\n",
    "track_masks = track_masks > .5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2210566/2931543681.py:32: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  proj = torch.tensor(proj, device=device).float()\n",
      "/tmp/ipykernel_2210566/2931543681.py:48: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  colors = torch.tensor(img.reshape(-1, 3)).cuda()\n"
     ]
    }
   ],
   "source": [
    "rr.init(\"depth optimization rerun\", recording_id=uuid.uuid4())\n",
    "rr.connect()\n",
    "\n",
    "rr.log(\"world\", rr.Clear(recursive=True))\n",
    "rr.log(\"log\", rr.Clear(recursive=True))\n",
    "rr.log(\"world\", rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, static=True)\n",
    "rr.log(\"world/scene\", rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, static=True)\n",
    "\n",
    "# Dynamic continuous reconstruction\n",
    "\n",
    "visualize_video_rerun(\n",
    "    best_trajectory,\n",
    "    depths_dc_uni,\n",
    "    imgs,\n",
    "    proj,\n",
    "    uncertainties=uncertainties_1,\n",
    "    cam_every=18,\n",
    "    depth_every=1,\n",
    "    subsample_pts=5, # 5\n",
    "    subsample_dynamic_pts=1, # 2\n",
    "    uncertainty_thresh=uncert_thresh,\n",
    "    track_masks=track_masks,\n",
    "    max_depth_masks=max_depth_masks.to(track_masks.device),\n",
    "    follow_cam=\"world\",\n",
    "    radii=2,\n",
    "    filter_depth_threshold=0.05,\n",
    "    step=0,\n",
    "    static_points_accumulate=True,\n",
    "    highlight_dynamic=0.01,\n",
    "    image_plane_distance=0.1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2210566/2931543681.py:32: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  proj = torch.tensor(proj, device=device).float()\n",
      "/tmp/ipykernel_2210566/2931543681.py:48: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  colors = torch.tensor(img.reshape(-1, 3)).cuda()\n"
     ]
    }
   ],
   "source": [
    "rr.init(\"depth optimization rerun\", recording_id=uuid.uuid4())\n",
    "rr.connect()\n",
    "\n",
    "rr.log(\"world\", rr.Clear(recursive=True))\n",
    "rr.log(\"log\", rr.Clear(recursive=True))\n",
    "rr.log(\"world\", rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, static=True)\n",
    "rr.log(\"world/scene\", rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, static=True)\n",
    "\n",
    "# Dynamic discrete reconstruction\n",
    "\n",
    "visualize_video_rerun(\n",
    "    best_trajectory,\n",
    "    depths_dc_uni,\n",
    "    imgs,\n",
    "    proj,\n",
    "    uncertainties=uncertainties_1,\n",
    "    cam_every=18,\n",
    "    depth_every=18,\n",
    "    subsample_pts=5, # 5\n",
    "    subsample_dynamic_pts=1, # 2\n",
    "    uncertainty_thresh=uncert_thresh,\n",
    "    max_depth_masks=max_depth_masks.to(track_masks.device),\n",
    "    follow_cam=\"world\",\n",
    "    radii=2,\n",
    "    filter_depth_threshold=0.05,\n",
    "    step=0,\n",
    "    static_points_accumulate=True,\n",
    "    highlight_dynamic=0.01,\n",
    "    image_plane_distance=0.1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MoviePy - Building video media/anycam/4D/running_dog_pexels/imgs/video.mp4.\n",
      "MoviePy - Writing video media/anycam/4D/running_dog_pexels/imgs/video.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                         \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MoviePy - Done !\n",
      "MoviePy - video ready media/anycam/4D/running_dog_pexels/imgs/video.mp4\n",
      "MoviePy - Building video media/anycam/4D/running_dog_pexels/depths/video.mp4.\n",
      "MoviePy - Writing video media/anycam/4D/running_dog_pexels/depths/video.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                         \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MoviePy - Done !\n",
      "MoviePy - video ready media/anycam/4D/running_dog_pexels/depths/video.mp4\n",
      "MoviePy - Building video media/anycam/4D/running_dog_pexels/uncertainties/video.mp4.\n",
      "MoviePy - Writing video media/anycam/4D/running_dog_pexels/uncertainties/video.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                         \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MoviePy - Done !\n",
      "MoviePy - video ready media/anycam/4D/running_dog_pexels/uncertainties/video.mp4\n",
      "MoviePy - Building video media/anycam/4D/running_dog_pexels/flows_occs_fwd/video.mp4.\n",
      "MoviePy - Writing video media/anycam/4D/running_dog_pexels/flows_occs_fwd/video.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                         \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MoviePy - Done !\n",
      "MoviePy - video ready media/anycam/4D/running_dog_pexels/flows_occs_fwd/video.mp4\n"
     ]
    }
   ],
   "source": [
    "# Export to input images, depth, flow, and uncertainties to video and frames\n",
    "\n",
    "import cv2\n",
    "from moviepy import VideoClip\n",
    "\n",
    "def save_to_video_and_frames(frames, out_path, fps=30):\n",
    "    out_path.mkdir(exist_ok=True, parents=True)\n",
    "    frames_dir = out_path / \"frames\"\n",
    "    frames_dir.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "    for i, frame in enumerate(frames):\n",
    "        cv2.imwrite(str(frames_dir / f\"{i:04d}.png\"), cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))\n",
    "\n",
    "    clip = VideoClip(lambda t: frames[int(t * fps)], duration=len(frames) / fps)\n",
    "    clip.write_videofile(str(out_path / \"video.mp4\"), codec=\"libx264\", fps=fps)\n",
    "\n",
    "def save_all(out_path, imgs, depths, uncertainties, flows_occs_fwd, fps=30):\n",
    "    save_to_video_and_frames([(img * 255).astype(np.uint8) for img in imgs], out_path / \"imgs\", fps)\n",
    "\n",
    "    depths_c = [(color_tensor(((1 / (depths[i, 0] * 10)) * 5).clamp(0, 1), cmap=\"plasma\", norm=False) * 255).numpy().astype(np.uint8) for i in range(len(depths))]\n",
    "    save_to_video_and_frames(depths_c, out_path / \"depths\", fps)\n",
    "\n",
    "    uncertainties_c = [(color_tensor((uncertainties[i, 0] / 0.08).clamp(0, 1), cmap=\"plasma\", norm=False) * 255).numpy().astype(np.uint8) for i in range(len(uncertainties))]\n",
    "    save_to_video_and_frames(uncertainties_c, out_path / \"uncertainties\", fps)\n",
    "\n",
    "    flows = flow_to_image(flows_occs_fwd[:, :2].cpu()).permute(0, 2, 3, 1).numpy()\n",
    "    save_to_video_and_frames(flows, out_path / \"flows_occs_fwd\", fps)\n",
    "\n",
    "save_all(Path(\"media/anycam/4D/running_dog_pexels\"), imgs, depths_dc_uni, uncertainties, seq_flows_occs_fwd, fps=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
